#!/usr/bin/env python3

import pandas as pd
from qiutil import DEFAULT_END_DATE, DEFAULT_GOLD_FILE, DEFAULT_START_DATE

# Instead of using classes
# Dataframe is easier to understand
# Portfolio can be represented in dataframe
# INDEX: (PTF_NAME, START, END) VALUE: STOCK
class Portfolio:
    def __init__(self, start, end, stocks, mrt):
        self.start = start
        self.end = end
        self.stocks = stocks
        self.mrt_in_span = mrt.query(f"PREV_MON >= '{self.start}' & PREV_MON <= '{self.end}'")
        self.df = pd.merge(self.stocks, self.mrt_in_span, on=['STOCK_CODE'])

    def overall_return(self):
        # 建立观察期空 DataFrame
        date_range = pd.date_range(start=self.start, end=self.end, freq='M').strftime("%Y%m%d")
        mr = pd.DataFrame({'DATE':date_range})
        mr['EQUAL_RETURN'] = 0
        mr['VALUE_RETURN'] = 0
        mr['EQUAL_FACTOR'] = 1
        mr['VALUE_FACTOR'] = 1
        mr['EQUAL_VALUE'] = 1
        mr['VALUE_VALUE'] = 1

        # 计算 等权重、市值加权 组合收益率
        self.df['EQUAL_RETURN'] = self.df['EQUAL_WEIGHT'] * self.df['MONTHLY_RETURN']
        self.df['VALUE_RETURN'] = self.df['VALUE_WEIGHT'] * self.df['MONTHLY_RETURN']

        for dt in date_range:
            dt_index = mr[mr['DATE']==dt].index
            mr.loc[dt_index,'EQUAL_RETURN'] = self.df[self.df['PREV_MON']==dt]['EQUAL_RETURN'].sum()
            mr.loc[dt_index,'VALUE_RETURN'] = self.df[self.df['PREV_MON']==dt]['VALUE_RETURN'].sum()
            mr.loc[dt_index,'EQUAL_FACTOR'] = 1 + mr.loc[dt_index,'EQUAL_RETURN']
            mr.loc[dt_index,'VALUE_FACTOR'] = 1 + mr.loc[dt_index, 'VALUE_RETURN']
            mr.loc[dt_index,'EQUAL_VALUE'] = mr[mr['DATE']<=dt]['EQUAL_FACTOR'].prod()
            mr.loc[dt_index,'VALUE_VALUE'] = mr[mr['DATE']<=dt]['VALUE_FACTOR'].prod()

        mr.drop(labels=['EQUAL_FACTOR','VALUE_FACTOR'],axis=1,inplace=True)
        return mr.copy()

def create_portfolios(gd, start=DEFAULT_START_DATE, end=DEFAULT_END_DATE):
    ptfs = pd.DataFrame()
    periods = gd.df.index.get_level_values(0).drop_duplicates()
    for cs_date in periods:
        df = gd.df.loc[cs_date].sort_values(by='ADJ_ROE',ascending=False).iloc[0:100,]
        df = df[['STOCK_CODE','CS_DATE','ADJ_ROE', 'FLOAT_MV']]
        # 添加 等权重 标记
        df['EQUAL_WEIGHT'] = 1 / 100

        # 添加 市值权重 标记
        float_mv_sum = df['FLOAT_MV'].sum()
        df['VALUE_WEIGHT'] = df['FLOAT_MV'] / float_mv_sum
        df["PTF_NAME"] = "TOP100_" + cs_date
        df['PTF_START_DATE'] = start
        df['PTF_END_DATE'] = end
        ptfs = ptfs.append(df, ignore_index=True)
    return ptfs.copy()

def save_portfolios(df, data_file=DEFAULT_GOLD_FILE):
    df.to_hdf(data_file, "ptfs")
